import numpy as np
import os
import struct

MAX_STRING_LEN = 256
HF_MAX_STR_LEN = 128
QT_MAX_STR_LEN = 64


class NameLoc:
    def __init__(self):
        self.type = "u"
        self.length = 0
        self.on_grid = False
        self.time_type = ""
        self.time = 0.0
        self.Max = 0.0
        self.Min = 0.0
        self.Avg = 0.0
        self.name = ""
        self.unit = ""
        self.fcount = 0
        self.offset = 0

class MeshReader:
    def __init__(self):
        self.clear()

    def clear(self):
        self.casename = ""
        self.stat_props = []  # Static property list
        self.dym_props = []  # Dynamic property list
        self.Gross2Net = np.array([], dtype=np.int32)
        self.Net2Gross = np.array([], dtype=np.int32)
        self.GridInfo = np.zeros(32, dtype=np.int32)
        self.AxisOffset = np.zeros(3, dtype=np.float64)
        self.origin_net_size = 0
        self.origin_GridInfo = np.zeros(32, dtype=np.int32)
        self.origin_Gross2Net = np.array([], dtype=np.int32)

    # 0-Cartesian, 1-Stage, 2-CPG, 3-PEBI
    def get_grid_type(self):
        if self.GridInfo[0] == 0:
            return "Cartesian"
        elif self.GridInfo[0] == 1:
            return "Stage"
        elif self.GridInfo[0] == 2:
            return "CPG"
        elif self.GridInfo[0] == 3:
            return "PEBI"

    def get_porolvl(self):
        return self.GridInfo[1]  # 双孔模型层数

    def get_grid_ni(self):
        return self.GridInfo[3]

    def get_grid_nj(self):
        return self.GridInfo[4]

    def get_grid_nk(self):
        return self.GridInfo[5]

    def get_net_size(self):
        return len(self.Net2Gross)  # GridInfo[8]

    def get_gross_size(self):
        return len(self.Gross2Net)  # GridInfo[2]

    def get_time_steps(self):
        return len(self.dym_props)

    def get_base_grid_num(self):
        return self.GridInfo[2] - self.GridInfo[6] - self.GridInfo[7]

    def get_active_base_grid_num(self):
        return self.GridInfo[9] + self.GridInfo[11] * (self.GridInfo[1] - 1)

    def get_fine_grid_num(self):
        return self.GridInfo[6]

    def get_active_fine_grid_num(self):
        return self.GridInfo[10] + self.GridInfo[12] * (self.GridInfo[1] - 1)

    def scan_bin_array(self, fp, a):
        """
        Read a binary array from the file pointer.

        :param fp: File pointer.
        :param a: NameLoc object.

        """
        a.offset = fp.tell()

        # Read char_len (UGINT64, 8 bytes, unsigned long long)
        char_len, a.time = struct.unpack("<q d", fp.read(16))
        a.type = fp.read(1).decode("latin1")

        # Read pname (HF_MAX_STR_LEN bytes)
        a.name = fp.read(HF_MAX_STR_LEN).split(b"\x00", 1)[0].decode("latin1")

        # Read punit (HF_MAX_STR_LEN bytes)
        punit_bytes = fp.read(HF_MAX_STR_LEN)
        a.unit = punit_bytes[:QT_MAX_STR_LEN].split(b"\x00", 1)[0].decode("latin1")

        offset = QT_MAX_STR_LEN
        a.Max, a.Min, a.Avg = struct.unpack("<ddd", punit_bytes[offset : offset + 24])
        offset += 24
        a.time_type = punit_bytes[offset : offset + 1].decode("latin1")
        a.on_grid = punit_bytes[offset + 1 : offset + 2] == b"y"

        net_len = char_len - 8 - 1 - MAX_STRING_LEN
        fp.seek(net_len, 1)

        if a.type == "d":
            a.length = net_len // 8
        elif a.type == "i":
            a.length = net_len // 4
        else:
            raise ValueError(f"Unknown data type: {a.type}")

    def get_static_data(self, name, prop_array, remap=False):
        """
        Get static data from the binary file.

        :param name: Property name.
        :param prop_array: Dictionary to store the property data.
        :param remap: Whether to remap the property data.

        :return: True if the property is found, False otherwise.

        """
        prop = next((p for p in self.stat_props if p.name == name), None)
        if not prop:
            return False

        filename = f"{self.casename}_geom.bin"
        try:
            with open(filename, "rb") as stat_prop_fp:
                # Read char_len (UGINT64, 8 bytes, unsigned long long)
                stat_prop_fp.seek(prop.offset, os.SEEK_SET)
                char_len, prop_array["time"] = struct.unpack(
                    "<q d", stat_prop_fp.read(16)
                )
                prop_array["type"] = stat_prop_fp.read(1).decode("latin1")

                # Read pname (HF_MAX_STR_LEN bytes)
                prop_array["name"] = (
                    stat_prop_fp.read(HF_MAX_STR_LEN)
                    .split(b"\x00", 1)[0]
                    .decode("latin1")
                )

                # Read punit (HF_MAX_STR_LEN bytes)
                punit_bytes = stat_prop_fp.read(HF_MAX_STR_LEN)
                prop_array["unit"] = (
                    punit_bytes[:QT_MAX_STR_LEN].split(b"\x00", 1)[0].decode("latin1")
                )

                # Extract Max, Min, Avg, time_type, and on_grid
                offset = QT_MAX_STR_LEN
                prop_array["Max"], prop_array["Min"], prop_array["Avg"] = struct.unpack(
                    "<ddd", punit_bytes[offset : offset + 24]
                )
                prop_array["time_type"] = punit_bytes[offset + 24 : offset + 25].decode(
                    "latin1"
                )
                prop_array["on_grid"] = punit_bytes[offset + 25 : offset + 26] == b"y"

                # Calculate net_len and read data
                net_len = char_len - 8 - 1 - MAX_STRING_LEN
                if prop_array["type"] == "d":
                    prop_array["length"] = net_len // 8
                    prop_array["data"] = np.frombuffer(
                        stat_prop_fp.read(net_len), dtype=np.float64
                    )
                elif prop_array["type"] == "i":
                    prop_array["length"] = net_len // 4
                    prop_array["data"] = np.frombuffer(
                        stat_prop_fp.read(net_len), dtype=np.int32
                    )
                else:
                    raise ValueError(f"Unknown data type: {prop_array['type']}")

                if np.isnan(prop_array["data"]).any():
                    print(f"Warning: NaN values found in property {name}")
                if np.isinf(prop_array["data"]).any():
                    print(f"Warning: Inf values found in property {name}")

                # Handle remapping if needed
                if remap and prop_array["on_grid"]:
                    if hasattr(self, "Moves") and self.Moves:
                        prop_array["data"] = self.re_map(
                            prop_array["data"],
                            self.NetSize(),
                            self.origin_net_size,
                            self.Moves,
                        )

                return True
        except FileNotFoundError:
            print(f"Error: File '{filename}' not found.")
            return False

    def get_dynamic_data(self, tstep, name, prop_array, remap=False):
        """
        Get dynamic data from the binary file.

        :param tstep: Time step.
        :param name: Property name.
        :param prop_array: Dictionary to store the property data.
        :param remap: Whether to remap the property data.

        :return: True if the property is found, False otherwise.

        """
        if tstep > len(self.dym_props):
            return False

        props = self.dym_props[tstep - 1]
        prop = next((p for p in props if p.name == name), None)
        if not prop:
            return False

        filename = f"{self.casename}_stu{prop.fcount if prop.fcount > 0 else ''}.bin"
        try:
            with open(filename, "rb") as dym_prop_fp:
                dym_prop_fp.seek(prop.offset, os.SEEK_SET)
                char_len, prop_array["time"] = struct.unpack(
                    "<q d", dym_prop_fp.read(16)
                )
                prop_array["type"] = dym_prop_fp.read(1).decode("latin1")

                # Read pname (HF_MAX_STR_LEN bytes)
                prop_array["name"] = (
                    dym_prop_fp.read(HF_MAX_STR_LEN)
                    .split(b"\x00", 1)[0]
                    .decode("latin1")
                )

                # Read punit (HF_MAX_STR_LEN bytes)
                punit_bytes = dym_prop_fp.read(HF_MAX_STR_LEN)
                prop_array["unit"] = (
                    punit_bytes[:QT_MAX_STR_LEN].split(b"\x00", 1)[0].decode("latin1")
                )

                # Extract Max, Min, Avg, time_type, and on_grid
                offset = QT_MAX_STR_LEN
                prop_array["Max"], prop_array["Min"], prop_array["Avg"] = struct.unpack(
                    "<ddd", punit_bytes[offset : offset + 24]
                )
                prop_array["time_type"] = punit_bytes[offset + 24 : offset + 25].decode(
                    "latin1"
                )
                prop_array["on_grid"] = punit_bytes[offset + 25 : offset + 26] == b"y"

                # Calculate net_len and read data
                net_len = char_len - 8 - 1 - MAX_STRING_LEN
                if prop_array["type"] == "d":
                    prop_array["length"] = net_len // 8
                    prop_array["data"] = np.frombuffer(
                        dym_prop_fp.read(net_len), dtype=np.float64
                    )
                elif prop_array["type"] == "i":
                    prop_array["length"] = net_len // 4
                    prop_array["data"] = np.frombuffer(
                        dym_prop_fp.read(net_len), dtype=np.int32
                    )
                else:
                    raise ValueError(f"Unknown data type: {prop_array['type']}")

                if np.isnan(prop_array["data"]).any():
                    print(f"Warning: NaN values found in property {name}")
                if np.isinf(prop_array["data"]).any():
                    print(f"Warning: Inf values found in property {name}")

                # Handle remapping if needed
                if remap and prop_array["on_grid"]:
                    if hasattr(self, "Moves") and self.Moves:
                        prop_array["data"] = self.re_map(
                            prop_array["data"],
                            self.NetSize(),
                            self.origin_net_size,
                            self.Moves,
                        )

                return True
        except FileNotFoundError:
            print(f"Error: File '{filename}' not found.")
            return False

    def scan(self, casename):
        """
        Scan the binary files for the specified case.

        :param casename: Case name.

        """
        self.clear()
        self.casename = casename

        # Read dynamic properties
        print(f"Reading dynamic properties for case {casename}")
        fcount = 0
        while True:
            filename = f"{self.casename}_stu{fcount if fcount != 0 else ''}.bin"
            if not os.path.exists(filename):
                if fcount == 0:
                    print(f'File "{filename}" does not exist!')
                break

            with open(filename, "rb") as fp:
                cur_time = -float("inf")
                while True:
                    current_position = fp.tell()
                    if not fp.read(1):
                        break
                    fp.seek(current_position, os.SEEK_SET)

                    tmp_prop = NameLoc()
                    tmp_prop.fcount = fcount
                    self.scan_bin_array(fp, tmp_prop)

                    if tmp_prop.time != cur_time or not self.dym_props:
                        step = len(self.dym_props)
                        for s in range(len(self.dym_props)):
                            if (
                                self.dym_props[s]
                                and self.dym_props[s][0].time == tmp_prop.time
                            ):
                                step = s
                                break

                        if step < len(self.dym_props):
                            self.dym_props[step].append(tmp_prop)
                        else:
                            self.dym_props.append([tmp_prop])

                        cur_time = tmp_prop.time
                    else:
                        self.dym_props[-1].append(tmp_prop)
            fcount += 1

        # Read static properties
        filename = f"{self.casename}_geom.bin"
        print(f"Reading static properties from {filename}")
        if os.path.exists(filename):
            with open(filename, "rb") as stat_prop_fp:
                while True:
                    current_position = stat_prop_fp.tell()
                    if not stat_prop_fp.read(1):
                        break
                    stat_prop_fp.seek(current_position, os.SEEK_SET)

                    tmp_prop = NameLoc()
                    self.scan_bin_array(stat_prop_fp, tmp_prop)

                    if tmp_prop.name == "G2NMAP":
                        gross_size = tmp_prop.length
                        stat_prop_fp.seek(
                            tmp_prop.offset + 17 + MAX_STRING_LEN, os.SEEK_SET
                        )
                        data_bytes = stat_prop_fp.read(gross_size * 4)
                        self.Gross2Net = np.frombuffer(data_bytes, dtype=np.int32)

                    if tmp_prop.name == "N2GMAP":
                        net_size = tmp_prop.length
                        stat_prop_fp.seek(
                            tmp_prop.offset + 17 + MAX_STRING_LEN, os.SEEK_SET
                        )
                        data_bytes = stat_prop_fp.read(net_size * 4)
                        self.Net2Gross = np.frombuffer(data_bytes, dtype=np.int32)

                    self.stat_props.append(tmp_prop)

        # Get GRIDINFO
        GRIDINFO_array = {}
        if self.get_static_data("GRIDINFO", GRIDINFO_array):
            if GRIDINFO_array["length"] == 32:
                self.GridInfo[:20] = GRIDINFO_array["data"][:20].astype(np.int32)
                self.AxisOffset = GRIDINFO_array["data"][20:23]
            else:
                print("Error: GRIDINFO length is not 32")
        else:
            print("Error: GRIDINFO data not found")

        # Consistency check
        if self.GridInfo[2] != len(self.Gross2Net):
            print("Error: G2NMAP length does not match GRIDINFO[2]")

        if self.GridInfo[8] != len(self.Net2Gross):
            print("Error: N2GMAP length does not match GRIDINFO[8]")

        self.origin_net_size = len(self.Net2Gross)
        self.origin_GridInfo = self.GridInfo.copy()
        self.origin_Gross2Net = self.Gross2Net.copy()

    def print_grid_info(self):
        print(
            f"Overall grid size = {self.get_gross_size()}, active grid = {self.get_net_size()}, number of time steps = {self.get_time_steps()}"
        )
        print(f"Axis_offset = {tuple(self.AxisOffset)}")
        print(
            f"NI = {self.get_grid_ni()}, NJ = {self.get_grid_nj()}, NK = {self.get_grid_nk()}, NLayer = {self.get_grid_nk() // self.get_porolvl()}"
        )
        print(
            f"PoroLvl = {self.get_porolvl()}, number of visible elements = length(NODE_LIST)/8 = {self.get_net_size()}"
        )

    def print_grid_nums(self):
        print(
            f"Base grid number = {self.get_base_grid_num()}, active base grid number = {self.get_active_base_grid_num()}"
        )
        print(
            f"Fine grid number = {self.get_fine_grid_num()}, active fine grid number = {self.get_active_fine_grid_num()}"
        )

    def print_static_props(self):
        print("Static properties:")
        for info in self.stat_props:
            print(
                f"    Name: {info.name}, length: {info.length}, unit: {info.unit}, is at cell center: {info.on_grid}"
            )

    def print_dynamic_props(self):
        print("Dynamic properties:")
        for dynamic_infos in self.dym_props:
            if dynamic_infos:
                print(
                    f"  Time: {dynamic_infos[0].time}, time type: {dynamic_infos[0].time_type}"
                )
            for info in dynamic_infos:
                print(
                    f"    Name: {info.name}, length: {info.length}, unit: {info.unit}, is at cell center: {info.on_grid}"
                )

class WellReader:
    def __init__(self):
        self.clear()

    def clear(self):
        self.casename = ""
        self.AxisOffset = np.zeros(3, dtype=np.float64)
        self.origin_WELL_size = 0
        self.WELL_num = []
        self.WELL_name = []
        self.WELL_node_count = []
        self.WELL_production_mode = []
        self.WELL_timestep = []
        self.WELL_stat_props = [] 
        self.WELL_dym_props = {}  
    
    def scan(self, casename):
        """
        Scan the out file and read the data.

        :param casename: _wstu.out file name.

        """
        self.clear()
        self.casename = casename

        # Read well properties
        filename = f"{self.casename}_wstu.out"
        print(f"Reading WELL properties from {filename}")
        if not os.path.exists(filename):
            print(f'File "{filename}" does not exist!')
        else:
            with open(filename, 'rb') as file:
                current_timestep = None
                current_well_properties = None
                current_well_num = None
                well_props_list = {}

                for line in file:
                    offset = file.tell() - len(line)
                    line = line.strip()
                    if not line:
                        continue
                    line_str = line.decode('utf-8')
                    fields = line_str.split()

                    if fields[0] == 'XCOORD0':
                        self.AxisOffset[0] = float(fields[1])
                    elif fields[0] == 'YCOORD0':
                        self.AxisOffset[1] = float(fields[1])
                    elif fields[0] == 'Time':
                        if current_well_properties:
                            well_props_list[current_well_num] = current_well_properties

                        if current_timestep is not None:
                            if current_timestep == 0.0:
                                self.WELL_stat_props =  {
                                    k: [(prop_name, unit, offset) for prop_name, unit, offset in v 
                                        if prop_name in ['XCOORD', 'YCOORD', 'DEPTH', 'STAGE', 'OUTLET']]
                                    for k, v in well_props_list.items()
                                }
                                self.WELL_dym_props[0.0] = {
                                    k: [(prop_name, unit, offset) for prop_name, unit, offset in v 
                                        if prop_name not in ['XCOORD', 'YCOORD', 'DEPTH', 'STAGE', 'OUTLET']]
                                    for k, v in well_props_list.items()
                                    if any(prop_name not in ['XCOORD', 'YCOORD', 'DEPTH', 'STAGE', 'OUTLET'] for prop_name, _, _ in v)
                                }
                            else:
                                self.WELL_dym_props[current_timestep] = well_props_list

                        current_timestep = float(fields[1])
                        well_props_list = {}
                        current_well_properties = []
                        self.WELL_timestep.append(current_timestep)

                    elif fields[0] == 'WELL':
                        if current_well_properties:
                            well_props_list[current_well_num] = current_well_properties
                            
                        current_well_num = int(fields[1])
                        current_well_properties = []

                        if(current_timestep == 0.0):
                            self.WELL_num.append(current_well_num)
                            self.WELL_name.append(fields[2].strip("'"))
                            self.WELL_node_count.append(int(fields[3]))
                            self.WELL_production_mode.append(' '.join(fields[4:]).strip("'"))

                    elif current_well_properties is not None:
                        property_name = fields[0]
                        unit = fields[1]
                        current_well_properties.append((property_name, unit, offset))

                # 保存最后一个井的属性
                if current_well_properties is not None:
                    well_props_list[current_well_num] = current_well_properties

                if current_timestep is not None:
                    if current_timestep == 0.0:
                        self.WELL_stat_props = {
                            k: [(prop_name, unit, offset) for prop_name, unit, offset in v 
                                if prop_name in ['XCOORD', 'YCOORD', 'DEPTH', 'STAGE', 'OUTLET']]
                            for k, v in well_props_list.items()
                        }
                        self.WELL_dym_props[0.0] = {
                            k: [(prop_name, unit, offset) for prop_name, unit, offset in v 
                                if prop_name not in ['XCOORD', 'YCOORD', 'DEPTH', 'STAGE', 'OUTLET']]
                            for k, v in well_props_list.items()
                            if any(prop_name not in ['XCOORD', 'YCOORD', 'DEPTH', 'STAGE', 'OUTLET'] for prop_name, _, _ in v)
                        }
                    else :
                        self.WELL_dym_props[current_timestep] = well_props_list

    def print_well_info(self):
        print("Well information:")
        print(
            f"Number of well = {len(self.WELL_num)}, Number of time steps = {len(self.WELL_timestep)}"
        )
        print(f"Axis_offset = {tuple(self.AxisOffset)}")

    def print_static_props(self):
        print("Static properties:")
        for well_num in self.WELL_num:
            if well_num in self.WELL_stat_props:
                print(f"  WELL {well_num}:")
                for property_name, unit, offset in self.WELL_stat_props[well_num]:
                    index = self.WELL_num.index(well_num)
                    print(f"    Name: {property_name}, length: {self.WELL_node_count[index]}, unit: {unit}, Offset: {offset}")
            else:
                print(f"Well number {well_num} is not found in static properties.")

    def print_dynamic_props(self):
        print("Dynamic properties:")
        for well_num in self.WELL_num:
            print(f" WELL {well_num}:")
            for timestep, wells in self.WELL_dym_props.items():
                print(f"  Time step: {timestep}")
                if well_num in wells:
                    for property_name, unit, offset in wells[well_num]:
                        index = self.WELL_num.index(well_num)
                        print(f"    Property Name: {property_name}, length: {self.WELL_node_count[index]}, unit: {unit}, Offset: {offset}")
                else:
                    print(f"  Well number {well_num} is not found in dynamic properties for this timestep.")
    
    def get_static_data(self, well_num, name, prop_array):
        if well_num in self.WELL_stat_props:
            for prop_name, unit, offset in self.WELL_stat_props[well_num]:
                if prop_name == name:
                    filename = f"{self.casename}_wstu.out"
                    if not os.path.exists(filename):
                        print(f'File "{filename}" does not exist!')
                        return None
                    with open(filename, 'rb') as file:
                        file.seek(offset)
                        line = file.readline().strip()

                        parts = line.split()
                        unit = parts[1]
                        data_values = parts[2:]

                        if unit.lower() == "integer":
                            data = list(map(int, data_values))
                        else:
                            data = list(map(float, data_values))
                        prop_array[name] = data
                    return offset
        return None

    def get_dynamic_data(self, tstep, well_num, name, prop_array):
        if tstep in self.WELL_dym_props and well_num in self.WELL_dym_props[tstep]:
            for prop_name, unit, offset in self.WELL_dym_props[tstep][well_num]:
                if prop_name == name:
                    filename = f"{self.casename}_wstu.out"
                    if not os.path.exists(filename):
                        print(f'File "{filename}" does not exist!')
                        return None
                    with open(filename, 'rb') as file:
                        file.seek(offset)
                        line = file.readline().strip()

                        parts = line.split()
                        unit = parts[1]
                        data_values = parts[2:]

                        if unit.lower() == "integer":
                            data = list(map(int, data_values))
                        else:
                            data = list(map(float, data_values))
                        prop_array[name] = data
                    return offset
        return None
